"""
风险管理模块
"""
from typing import Dict, Optional
from dataclasses import dataclass
from src.utils.logger import get_logger

logger = get_logger(__name__)


@dataclass
class RiskControl:
    """风险控制参数"""
    stop_loss_pct: float  # 单笔止损 (%)
    take_profit_pct: float  # 单笔止盈 (%)
    max_positions: int  # 最大并发持仓数
    max_daily_loss_pct: float  # 最大日亏 (%)
    max_drawdown_pct: float  # 最大回撤 (%)
    max_position_bars: int  # 最大持仓K线数


class RiskManager:
    """风险管理器"""

    def __init__(self, risk_control: RiskControl, initial_balance: float):
        """初始化"""
        self.risk = risk_control
        self.initial_balance = initial_balance
        self.current_balance = initial_balance
        self.peak_balance = initial_balance
        self.daily_loss = 0.0
        self.positions = []

    def update_balance(self, new_balance: float):
        """更新账户余额"""
        self.current_balance = new_balance
        if new_balance > self.peak_balance:
            self.peak_balance = new_balance

    def can_open_position(self, position_size: float) -> bool:
        """
        检查是否可以开仓

        Args:
            position_size: 拟开仓的持仓规模

        Returns:
            是否允许开仓
        """
        # 检查并发持仓数
        if len(self.positions) >= self.risk.max_positions:
            logger.warning(f"已达到最大并发持仓数: {self.risk.max_positions}")
            return False

        # 检查账户是否有足够资金
        if position_size > self.current_balance * 0.9:  # 保留10%的缓冲
            logger.warning(f"账户余额不足")
            return False

        # 检查日亏损限制
        if self.daily_loss >= self.initial_balance * (self.risk.max_daily_loss_pct / 100):
            logger.warning(f"已达到最大日亏损: {self.risk.max_daily_loss_pct}%")
            return False

        # 检查最大回撤
        drawdown_pct = ((self.peak_balance - self.current_balance) / self.peak_balance) * 100
        if drawdown_pct >= self.risk.max_drawdown_pct:
            logger.warning(f"已达到最大回撤: {drawdown_pct:.2f}%")
            return False

        return True

    def calculate_position_size(self, entry_price: float, stop_loss_price: float) -> float:
        """
        根据风险计算仓位大小

        Args:
            entry_price: 入场价格
            stop_loss_price: 止损价格

        Returns:
            推荐仓位大小
        """
        # 单笔最大风险
        risk_amount = self.current_balance * 0.02  # 每笔风险2%账户

        # 止损距离
        stop_loss_distance = abs(entry_price - stop_loss_price)

        if stop_loss_distance == 0:
            return 0

        # 仓位大小 = 风险金额 / 止损距离
        position_size = risk_amount / stop_loss_distance

        return position_size

    def should_close_position(self, position_entry_bar: int, current_bar: int) -> bool:
        """
        检查是否应该平仓 (基于K线数量限制)

        Args:
            position_entry_bar: 开仓时的K线索引
            current_bar: 当前K线索引

        Returns:
            是否应该平仓
        """
        bars_held = current_bar - position_entry_bar

        if bars_held >= self.risk.max_position_bars:
            logger.info(f"已达到最大持仓K线数: {bars_held} >= {self.risk.max_position_bars}")
            return True

        return False
